import sys
import os

sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
import re
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import Dataset, DataLoader
import torch.optim as optim
import matplotlib.pyplot as plt
from tqdm import tqdm
from transformer.transformer import *

plt.rcParams["font.sans-serif"] = ["SimHei"]  # Mac本字体改为: ['Arial Unicode MS']
plt.rcParams["axes.unicode_minus"] = False

# 设备选择, 我们可以选择在cuda或者cpu上运行你的代码
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# 如果是Mac M芯片, 写这个.
# device = 'mps'

# 1- 定义特殊标识和全局变量
SOS_token = 0  # 起始标识和索引
EOS_token = 1  # 结束标识和索引
MAX_LENGTH = 10  # 限定句子的最大长度为10，含标点在内
data_path = "data/eng-fra-v2.txt"  # 语料库路径


# 2- 数据清洗
def normalizeString(line):
    # 将内容全部转成小写；去除每个句子前后空白（空格、制表符、回车）字符串
    line = line.lower().strip()

    # 在标点符号（.!?）的前面加上空格
    line = re.sub(r"([.!?])", r" \1", line)

    # 除了26个字母和基础的标点符号（.!?）这些以外的其他内容，全部清除，替换为空格
    line = re.sub(r"[^a-z.!?]+", " ", line)

    return line


# 3- 读取文件内容，处理得到字典
def my_getdata():
    # 1. 读取文件
    # 下面的写法更优，能够及时释放文件资源
    with open(data_path, mode="r", encoding="UTF-8") as file_obj:
        lines = file_obj.readlines()

    # 2. 对每一行都进行处理，得到 英文和法文 的句子对
    """
        目前的lines内容：["英语句子\t法语句子\n", "英语句子\t法语句子\n"]
        最终想要的结构是：[["英语句子","法语句子"]]
    """

    my_pairs = [
        [normalizeString(eng_fra) for eng_fra in line.split("\t")] for line in lines
    ]
    # print(my_pairs)

    # 3. 分别构建得到英文、法文的 word->index的映射关系
    # 3.1- 定义字典
    english_word2index = {"SOS": SOS_token, "EOS": EOS_token}
    english_word_n = 2
    france_word2index = {"SOS": SOS_token, "EOS": EOS_token}
    france_word_n = 2

    # 3.2- 对每个英语、法语句子进行分词处理；去重然后放到字典中
    for eng_fra in my_pairs:
        # 处理英语
        for word in eng_fra[0].split(" "):
            if word not in english_word2index:
                english_word2index[word] = english_word_n
                english_word_n += 1

        # 处理法语
        for word in eng_fra[1].split(" "):
            if word not in france_word2index:
                france_word2index[word] = france_word_n
                france_word_n += 1

    # 4. 分别构建得到英文、法文的 index->word的映射关系
    english_index2word = {index: word for word, index in english_word2index.items()}
    france_index2word = {index: word for word, index in france_word2index.items()}

    return (
        english_word2index,
        english_index2word,
        english_word_n,
        france_word2index,
        france_index2word,
        france_word_n,
        my_pairs,
    )


# 调用my_getdata方法，将返回值设置为全局变量。因为 后面都要用到
(
    english_word2index,
    english_index2word,
    english_word_n,
    france_word2index,
    france_index2word,
    france_word_n,
    my_pairs,
) = my_getdata()


# 4- 自定义Dataset
class MyPairsDataset(Dataset):
    def __init__(self, my_pairs):
        self.my_pairs = my_pairs  # 英语、法语句子对。嵌套列表
        self.sample_len = len(my_pairs)  # 样本条数

    def __len__(self):
        # 返回样本总条数
        return self.sample_len

    def __getitem__(self, index):
        # 1- index索引值调整，防止越界并且不考虑负索引
        index = min(max(index, 0), self.sample_len - 1)

        # 2- 获得字符串形式的 特征数据（英语句子） 和 目标值数据（法语句子）
        x = self.my_pairs[index][0]
        y = self.my_pairs[index][1]

        # 3- 句子中的词替换成对应的索引：先分词得到单词，通过单词查询得到对应的索引
        """
            为什么这里只是在句子的后面专门添加了句子的结束标识EOS_token，而没有添加句子的开始标识SOS_token？
            在Seq2Seq+注意力机制中，不管是编码器还是解码器中输入进去的句子的末尾都必须加上EOS_token，明确告诉处理已经结束了
            而SOS_token不是必须加的，可以在解码器中加
        """
        x = [english_word2index[word] for word in x.split(" ")]
        x.append(EOS_token)
        tensor_x = torch.tensor(x, dtype=torch.long, device=device)

        y = [france_word2index[word] for word in y.split(" ")]
        y.append(EOS_token)
        tensor_y = torch.tensor(y, dtype=torch.long, device=device)

        return tensor_x, tensor_y


# 5- 测试Dataset
def get_dataloader():
    # 1- 创建Dataset
    dataset = MyPairsDataset(my_pairs)

    # 2- 创建Dataloader
    torch.manual_seed(113)
    dataloader = DataLoader(dataset, batch_size=1, shuffle=True)

    return dataloader


def use_transformer(test_mode=False):
    # 常量部分
    d_model = 512
    dropout = 0.1
    head = 8
    clip = 1.0  # 梯度裁剪阈值
    print_every = 50  # 每多少个批次打印一次损失
    max_batches = 500 if test_mode else None  # 测试模式下只处理500个批次

    # 准备数据
    dataloader = get_dataloader()

    # 编码器部分
    en_embed = Embedding(vocab_size=english_word_n, d_model=d_model)
    en_pos = PositionalEncoding(dropout=dropout, d_model=d_model, max_len=60)
    en_multi_self_attn = MultiHeadAttention(
        d_model=d_model, head=head, dropout_p=dropout
    )
    en_ff = FeedForward(d_model=d_model, d_ff=1024, dropout_p=dropout)
    encoder_layer = EncoderLayer(
        d_model=d_model,
        self_attn=en_multi_self_attn,
        feed_forward=en_ff,
        dropout=dropout,
    )
    encoder = Encoder(layer=encoder_layer, N=6)

    # 解码器部分
    de_embed = Embedding(vocab_size=france_word_n, d_model=d_model)
    de_pos = PositionalEncoding(dropout=dropout, d_model=d_model, max_len=60)

    (
        de_multi_self_attn,
        de_multi_attn,
    ) = clones(
        MultiHeadAttention(d_model=d_model, head=head, dropout_p=dropout), obj_num=2
    )
    de_ff = FeedForward(d_model=d_model, d_ff=1024, dropout_p=dropout)
    decoder_layer = DecoderLayer(
        d_model=d_model,
        mask_self_attn=de_multi_self_attn,
        src_attn=de_multi_attn,
        feed_forward=de_ff,
        dropout=dropout,
    )
    decoder = Decoder(layer=decoder_layer, N=6)

    # 输出
    generator = Generator(d_model=d_model, vocab_size=france_word_n)

    # 组装得到Transformer框架
    my_transformer = MyTransformer(
        en_embed=nn.Sequential(en_embed, en_pos),
        encoder=encoder,
        de_embed=nn.Sequential(de_embed, de_pos),
        decoder=decoder,
        generator=generator,
    )
    # 将模型移到正确的设备上
    my_transformer = my_transformer.to(device)

    # 初始化模型参数
    for p in my_transformer.parameters():
        if p.dim() > 1:
            nn.init.xavier_uniform_(p)

    # 优化器初始化 - 使用单个优化器优化整个模型
    optimizer = optim.Adam(
        my_transformer.parameters(), lr=1e-4, betas=(0.9, 0.98), eps=1e-9
    )

    # 损失函数初始化 - 使用CrossEntropyLoss处理分类问题
    criterion = nn.CrossEntropyLoss(ignore_index=0)  # 忽略填充符的索引

    # 训练参数初始化
    epochs = 5
    plot_loss_list = []  # 记录损失值，用来绘制图形

    # 设置模型为训练模式
    my_transformer.train()

    # 循环训练
    for epoch in range(epochs):
        # 变量初始化
        print_loss_total = 0.0  # 用于打印的总损失
        plot_loss_total = 0.0  # 用于绘制图形的总损失

        print(f"\n训练轮次: {epoch+1}/{epochs}")

        for item, (x, y) in enumerate(tqdm(dataloader), start=1):
            # 梯度清零
            optimizer.zero_grad()

            # 确保输入和目标张量在正确的设备上
            x = x.to(device)
            y = y.to(device)

            # 根据实际输入序列长度动态创建掩码
            batch_size = x.shape[0]
            src_seq_len = x.shape[1]

            # 目标输入（去除最后一个token）和目标输出（去除第一个token）
            tgt_input = y[:, :-1]  # 解码器输入 (长度为原始长度-1)
            tgt_output = y[:, 1:]  # 预期输出

            # 使用解码器输入的实际长度创建掩码
            tgt_seq_len = tgt_input.shape[1]

            # 源掩码：创建全1矩阵表示所有位置都可见
            source_mask = torch.ones(
                size=(batch_size, 1, 1, src_seq_len), device=device
            )

            # 目标掩码：下三角矩阵，对角线及以下为1（可见）
            target_mask = torch.tril(
                torch.ones(size=(tgt_seq_len, tgt_seq_len), device=device), diagonal=0
            )
            target_mask = (
                target_mask.unsqueeze(0).unsqueeze(0).repeat(batch_size, 1, 1, 1)
            )

            # 前向传播
            output = my_transformer(x, tgt_input, source_mask, target_mask)

            # 计算损失
            # 重塑输出以适应CrossEntropyLoss的要求 [batch_size * seq_len, vocab_size]
            output = output.contiguous().view(-1, france_word_n)
            tgt_output = tgt_output.contiguous().view(-1)

            # 计算损失
            loss = criterion(output, tgt_output)

            # 反向传播
            loss.backward()

            # 梯度裁剪，防止梯度爆炸
            torch.nn.utils.clip_grad_norm_(my_transformer.parameters(), clip)

            # 参数更新
            optimizer.step()

            # 损失值累加
            print_loss_total += loss.item()
            plot_loss_total += loss.item()

            # 打印损失信息
            if item % print_every == 0:
                print_loss_avg = print_loss_total / print_every
                print(f"批次: {item}, 平均损失: {print_loss_avg:.4f}")
                print_loss_total = 0.0

            # 如果是测试模式，限制处理的批次数量
            if max_batches and item >= max_batches:
                print(f"\n测试模式：已处理{max_batches}个批次")
                break

        # 计算平均损失，考虑是否在测试模式下提前终止
        if max_batches and len(dataloader) > max_batches:
            epoch_loss_avg = plot_loss_total / max_batches
        else:
            epoch_loss_avg = plot_loss_total / len(dataloader)

        plot_loss_list.append(epoch_loss_avg)
        print(f"\nEpoch {epoch+1} 平均损失: {epoch_loss_avg:.4f}")

        # 如果是测试模式，只运行一个epoch
        if test_mode:
            print("\n测试模式：已完成一个epoch")
            break

    # 绘制损失曲线
    if plot_loss_list:
        plt.figure()
        plt.plot(range(1, len(plot_loss_list) + 1), plot_loss_list, marker="o")
        plt.title("Transformer 训练损失")
        plt.xlabel("Epochs")
        plt.ylabel("Loss")
        plt.grid(True)
        # 确保img文件夹存在
        import os

        os.makedirs("img", exist_ok=True)
        plt.savefig("img/transformer_loss.png")
        plt.show()
        print("\n损失曲线已保存到 img/transformer_loss.png")

    # 保存模型
    torch.save(my_transformer.state_dict(), "model/transformer_model.pkl")
    print("\n模型已保存到 model/transformer_model.pkl")

    return plot_loss_list
